import os
import matplotlib
import cv2
import numpy as np
from PIL import Image, ImageFont, ImageDraw

def __to_pil_img(img_buf):
    if isinstance(img_buf, bytes):
        img = np.asarray(bytearray(img_buf), dtype='uint8')
        image = cv2.imdecode(img, cv2.IMREAD_COLOR)
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    elif isinstance(img_buf, str) and os.path.exists(img_buf):
        image = Image.open(img_buf)
        image = image.convert('RGB')
    elif isinstance(img_buf, np.ndarray):
        image = Image.fromarray(cv2.cvtColor(img_buf, cv2.COLOR_BGR2RGB))
    elif isinstance(img_buf, Image.Image):
        image = img_buf
    else:
        image = None
    
    return image

class Canvas:
    def __init__(self, image, line_thickness=3, font=None):
        if isinstance(image, tuple):
            image = Image.new('RGBA', image[:2]) 
        else:
            image = __to_pil_img(image)
        self.draw = ImageDraw.Draw(image)
        self.line_thickness = line_thickness or max(min(image.size) // 200, 2)

        fontsize = max(round(max(image.size) / 40), 12)
        if isinstance(font, str) and os.path.isfile(font):
            self.font = ImageFont.truetype(font, fontsize, encoding='utf-8')
        else:
            font_path = font
            try:
                font = os.path.join(font_path, 'simsun.ttc')
                assert os.path.isfile(font)
            except:
                raise ValueError("Make sure simsun.ttc in %s" % font)
            self.font = ImageFont.truetype(font, fontsize, encoding='utf-8')
            
        
        self.dcolor = self.color_list()[0]
        self.image = image

    @staticmethod
    def color_list():
        # Return first 10 plt colors as (r,g,b) https://stackoverflow.com/questions/51350872/python-from-color-name-to-rgb
        def hex2rgb(h):
            return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))

        return [hex2rgb(h) for h in matplotlib.colors.TABLEAU_COLORS.values()]  # or BASE_ (8), CSS4_ (148), XKCD_ (949)
    
    def rect(self, box, color=None, label=None, line_thickness=None):
        color = color or self.dcolor
        line_thickness = line_thickness or self.line_thickness

        self.draw.rectangle(box, width=line_thickness, outline=tuple(color))

        if label:
            txt_width, txt_height = self.font.getsize(label)
            self.draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color))
            self.draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=self.font)

    def save(self, save_path):
        self.image.save(save_path)
    
    def show(self):
        self.image.show()